"""
CIFAR-10数据集演示
专门用于在CIFAR-10数据集上测试CLIP模型的零样本分类能力
"""

import torch
import torch.nn.functional as F
from torchvision import datasets, transforms
from PIL import Image
import matplotlib.pyplot as plt
import numpy as np
import tempfile
import os

from model import CLIPWrapper
from config import CLIP_CONFIG

def load_model():
    """加载CLIP模型"""
    model = CLIPWrapper(CLIP_CONFIG['model_name'])
    model.eval()
    return model

def load_cifar10_dataset():
    """加载CIFAR-10数据集"""
    # 数据预处理
    transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
    ])
    
    # 加载测试集
    test_dataset = datasets.CIFAR10(
        root='./data/cifar10', 
        train=False, 
        download=False,  # 从缓存加载
        transform=transform
    )
    
    return test_dataset

def zero_shot_classification(model, image_path, class_names):
    """零样本分类"""
    # 为每个类别创建描述
    texts = [f"a photo of a {name}" for name in class_names]
    
    # 加载图像
    image = Image.open(image_path).convert('RGB')
    
    # 处理输入
    inputs = model.processor(
        text=texts,
        images=image,
        return_tensors="pt",
        padding=True
    )
    
    # 编码
    with torch.no_grad():
        image_features = model.encode_image(inputs['pixel_values'])
        text_features = model.encode_text(inputs['input_ids'])
    
    # 计算相似度
    similarity = model.compute_similarity(image_features, text_features)
    
    # 转换为概率
    probs = F.softmax(similarity, dim=-1)
    probs = probs.squeeze().numpy()
    
    # 获取预测结果
    predicted_class_idx = np.argmax(probs)
    predicted_class = class_names[predicted_class_idx]
    confidence = probs[predicted_class_idx]
    
    return predicted_class, confidence, probs

def visualize_results(image_path, class_names, probs, true_label=None):
    """可视化分类结果"""
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))
    
    # 显示图像
    image = Image.open(image_path)
    ax1.imshow(image)
    ax1.axis('off')
    if true_label:
        ax1.set_title(f'True: {true_label}')
    else:
        ax1.set_title('Input Image')
    
    # 显示概率条形图
    y_pos = np.arange(len(class_names))
    colors = ['lightblue' for _ in class_names]
    
    # 高亮预测类别
    predicted_idx = np.argmax(probs)
    colors[predicted_idx] = 'lightcoral'
    
    ax2.barh(y_pos, probs, color=colors)
    ax2.set_yticks(y_pos)
    ax2.set_yticklabels(class_names)
    ax2.set_xlabel('Probability')
    ax2.set_title('Zero-Shot Classification Results')
    
    # 添加概率值
    for i, prob in enumerate(probs):
        ax2.text(prob + 0.01, i, f'{prob:.3f}', va='center')
    
    plt.tight_layout()
    plt.show()

def demo_zero_shot_classification():
    """演示零样本分类"""
    print("=== CIFAR-10零样本分类演示 ===")
    
    # 加载模型
    model = load_model()
    print("✓ 模型加载成功")
    
    # CIFAR-10类别名称
    class_names = ['airplane', 'automobile', 'bird', 'cat', 'deer', 
                   'dog', 'frog', 'horse', 'ship', 'truck']
    
    print(f"类别: {class_names}")
    
    try:
        # 加载CIFAR-10数据集
        test_dataset = load_cifar10_dataset()
        print(f"✓ 加载CIFAR-10测试集成功，样本数: {len(test_dataset)}")
        
        # 创建临时目录保存图像
        temp_dir = tempfile.mkdtemp(prefix='cifar_demo_')
        
        # 选择样本进行演示
        demo_samples = []
        samples_to_test = 10
        
        for i in range(samples_to_test):
            image, label = test_dataset[i]
            # 转换为PIL图像
            image_pil = transforms.ToPILImage()(image)
            
            # 保存到临时文件
            image_path = os.path.join(temp_dir, f"cifar10_sample_{i}.jpg")
            image_pil.save(image_path)
            
            demo_samples.append((image_path, class_names[label]))
        
        print(f"✓ 创建 {len(demo_samples)} 个演示样本")
        
        # 演示第一个样本
        print("\n=== 单个样本演示 ===")
        image_path, true_label = demo_samples[0]
        
        predicted_class, confidence, all_probs = zero_shot_classification(
            model, image_path, class_names
        )
        
        print(f"图像: CIFAR-10测试样本")
        print(f"真实类别: {true_label}")
        print(f"预测类别: {predicted_class}")
        print(f"置信度: {confidence:.4f}")
        print(f"预测正确: {predicted_class == true_label}")
        
        # 可视化结果
        visualize_results(image_path, class_names, all_probs, true_label)
        
        # 批量测试
        print("\n=== 批量测试结果 ===")
        correct_predictions = 0
        total_predictions = len(demo_samples)
        
        for i, (image_path, true_label) in enumerate(demo_samples):
            predicted_class, confidence, all_probs = zero_shot_classification(
                model, image_path, class_names
            )
            
            is_correct = predicted_class == true_label
            if is_correct:
                correct_predictions += 1
            
            status = "✓" if is_correct else "✗"
            print(f"{status} 样本 {i+1:2d}: 真实={true_label:10s} 预测={predicted_class:10s} 置信度={confidence:.4f}")
        
        accuracy = correct_predictions / total_predictions
        print(f"\n测试准确率: {accuracy:.2f} ({correct_predictions}/{total_predictions})")
        
        # 按类别统计准确率
        print("\n=== 按类别统计 ===")
        class_correct = {name: 0 for name in class_names}
        class_total = {name: 0 for name in class_names}
        
        for i, (image_path, true_label) in enumerate(demo_samples):
            predicted_class, confidence, all_probs = zero_shot_classification(
                model, image_path, class_names
            )
            
            class_total[true_label] += 1
            if predicted_class == true_label:
                class_correct[true_label] += 1
        
        for class_name in class_names:
            if class_total[class_name] > 0:
                class_acc = class_correct[class_name] / class_total[class_name]
                print(f"{class_name:10s}: {class_acc:.2f} ({class_correct[class_name]}/{class_total[class_name]})")
        
        print(f"\n临时文件保存在: {temp_dir}")
        
    except Exception as e:
        print(f"❌ 演示失败: {e}")
        print("请确保CIFAR-10数据集已下载到data/cifar10目录")

def demo_text_image_similarity():
    """演示文本-图像相似度"""
    print("\n=== CIFAR-10文本-图像相似度演示 ===")
    
    # 加载模型
    model = load_model()
    
    # CIFAR-10类别名称
    class_names = ['airplane', 'automobile', 'bird', 'cat', 'deer', 
                   'dog', 'frog', 'horse', 'ship', 'truck']
    
    try:
        # 加载CIFAR-10数据集
        test_dataset = load_cifar10_dataset()
        
        # 创建临时目录
        temp_dir = tempfile.mkdtemp(prefix='cifar_similarity_')
        
        # 选择样本
        demo_samples = []
        for i in range(5):
            image, label = test_dataset[i]
            image_pil = transforms.ToPILImage()(image)
            image_path = os.path.join(temp_dir, f"sample_{i}.jpg")
            image_pil.save(image_path)
            demo_samples.append((image_path, class_names[label]))
        
        # 测试文本查询
        test_queries = [
            "a photo of an airplane",
            "a photo of a cat",
            "a photo of a ship"
        ]
        
        for query in test_queries:
            print(f"\n查询: '{query}'")
            
            # 对每个样本计算相似度
            similarities = []
            for image_path, true_label in demo_samples:
                predicted_class, confidence, probs = zero_shot_classification(
                    model, image_path, class_names
                )
                
                # 找到查询对应的类别概率
                query_class = query.split()[-1]
                if query_class in class_names:
                    class_idx = class_names.index(query_class)
                    similarity = probs[class_idx]
                else:
                    similarity = 0.0
                
                similarities.append((true_label, similarity))
            
            # 按相似度排序
            similarities.sort(key=lambda x: x[1], reverse=True)
            
            print("相似度排名:")
            for i, (label, sim) in enumerate(similarities[:3]):  # 显示前3个
                print(f"  {i+1}. {label}: {sim:.4f}")
        
        print(f"\n临时文件保存在: {temp_dir}")
        
    except Exception as e:
        print(f"❌ 相似度演示失败: {e}")

def main():
    """主函数"""
    print("CIFAR-10数据集CLIP模型演示")
    print("=" * 50)
    
    # 演示零样本分类
    demo_zero_shot_classification()
    
    # 演示文本-图像相似度
    demo_text_image_similarity()
    
    print("\n演示完成!")

if __name__ == "__main__":
    main()